#include "crypto.h"
#include <string.h>
#include <time.h>

// PBKDF2 密钥派生函数
int pbkdf2_sha256(const uint8_t *password, size_t password_len,
                  const uint8_t *salt, size_t salt_len,
                  uint32_t iterations, uint8_t *key, size_t key_len) {
    if (!password || !salt || !key) return CRYPTO_ERROR_INVALID_PARAM;
    
    uint8_t u1[SHA256_DIGEST_SIZE];
    uint8_t u2[SHA256_DIGEST_SIZE];
    uint8_t block1[SHA256_DIGEST_SIZE];
    uint8_t block2[SHA256_DIGEST_SIZE];
    
    // 计算需要的块数
    size_t blocks = (key_len + SHA256_DIGEST_SIZE - 1) / SHA256_DIGEST_SIZE;
    
    for (size_t i = 1; i <= blocks; i++) {
        // 准备第一个块：salt || i (大端序)
        uint8_t block_input[salt_len + 4];
        memcpy(block_input, salt, salt_len);
        block_input[salt_len] = (i >> 24) & 0xFF;
        block_input[salt_len + 1] = (i >> 16) & 0xFF;
        block_input[salt_len + 2] = (i >> 8) & 0xFF;
        block_input[salt_len + 3] = i & 0xFF;
        
        // 计算 U1 = HMAC(password, salt || i)
        hmac_sha256(password, password_len, block_input, salt_len + 4, u1);
        memcpy(block1, u1, SHA256_DIGEST_SIZE);
        
        // 计算 U2, U3, ..., Uc
        for (uint32_t j = 1; j < iterations; j++) {
            hmac_sha256(password, password_len, u1, SHA256_DIGEST_SIZE, u2);
            memcpy(u1, u2, SHA256_DIGEST_SIZE);
            
            // 异或到结果中
            for (int k = 0; k < SHA256_DIGEST_SIZE; k++) {
                block1[k] ^= u1[k];
            }
        }
        
        // 复制到输出密钥
        size_t copy_len = (i == blocks) ? (key_len % SHA256_DIGEST_SIZE) : SHA256_DIGEST_SIZE;
        if (copy_len == 0) copy_len = SHA256_DIGEST_SIZE;
        
        memcpy(key + (i - 1) * SHA256_DIGEST_SIZE, block1, copy_len);
    }
    
    return CRYPTO_SUCCESS;
}

// HMAC-SHA256 实现
int hmac_sha256(const uint8_t *key, size_t key_len,
                const uint8_t *data, size_t data_len,
                uint8_t *mac) {
    if (!key || !data || !mac) return CRYPTO_ERROR_INVALID_PARAM;
    
    uint8_t k_pad[64];
    uint8_t o_key_pad[64];
    uint8_t i_key_pad[64];
    uint8_t inner_hash[SHA256_DIGEST_SIZE];
    
    // 如果密钥长度超过64字节，先进行哈希
    if (key_len > 64) {
        sha256_hash(key, key_len, k_pad);
        key_len = SHA256_DIGEST_SIZE;
    } else {
        memcpy(k_pad, key, key_len);
    }
    
    // 填充密钥到64字节
    memset(k_pad + key_len, 0, 64 - key_len);
    
    // 计算 o_key_pad 和 i_key_pad
    for (int i = 0; i < 64; i++) {
        o_key_pad[i] = k_pad[i] ^ 0x5c;
        i_key_pad[i] = k_pad[i] ^ 0x36;
    }
    
    // 计算内层哈希：H(K XOR ipad, text)
    sha256_ctx_t ctx;
    sha256_init(&ctx);
    sha256_update(&ctx, i_key_pad, 64);
    sha256_update(&ctx, data, data_len);
    sha256_final(&ctx, inner_hash);
    
    // 计算外层哈希：H(K XOR opad, inner_hash)
    sha256_init(&ctx);
    sha256_update(&ctx, o_key_pad, 64);
    sha256_update(&ctx, inner_hash, SHA256_DIGEST_SIZE);
    sha256_final(&ctx, mac);
    
    return CRYPTO_SUCCESS;
}

// 简单的随机数生成器（仅用于演示）
int crypto_random_bytes(uint8_t *buffer, size_t len) {
    if (!buffer) return CRYPTO_ERROR_INVALID_PARAM;
    
    // 使用当前时间作为种子
    static int initialized = 0;
    if (!initialized) {
        srand((unsigned int)time(NULL));
        initialized = 1;
    }
    
    // 生成随机字节
    for (size_t i = 0; i < len; i++) {
        buffer[i] = (uint8_t)(rand() % 256);
    }
    
    return CRYPTO_SUCCESS;
}

// 安全内存清零
void crypto_secure_zero(void *ptr, size_t len) {
    if (!ptr) return;
    
    volatile uint8_t *p = (volatile uint8_t *)ptr;
    while (len--) {
        *p++ = 0;
    }
}

// 安全的字符串比较（防止时序攻击）
int crypto_secure_compare(const uint8_t *a, const uint8_t *b, size_t len) {
    if (!a || !b) return 0;
    
    uint8_t result = 0;
    for (size_t i = 0; i < len; i++) {
        result |= a[i] ^ b[i];
    }
    
    return result == 0;
}

// 生成随机IV
int crypto_generate_iv(uint8_t *iv, size_t iv_len) {
    if (!iv) return CRYPTO_ERROR_INVALID_PARAM;
    
    return crypto_random_bytes(iv, iv_len);
}

// 生成随机盐值
int crypto_generate_salt(uint8_t *salt, size_t salt_len) {
    if (!salt) return CRYPTO_ERROR_INVALID_PARAM;
    
    return crypto_random_bytes(salt, salt_len);
}

// 计算数据的SHA256哈希值（便捷函数）
int crypto_hash_sha256(const char *data, uint8_t *digest) {
    if (!data || !digest) return CRYPTO_ERROR_INVALID_PARAM;
    
    return sha256_hash((const uint8_t *)data, strlen(data), digest);
}

// 将字节数组转换为十六进制字符串
void crypto_bytes_to_hex(const uint8_t *bytes, size_t len, char *hex) {
    if (!bytes || !hex) return;
    
    const char hex_chars[] = "0123456789abcdef";
    for (size_t i = 0; i < len; i++) {
        hex[i * 2] = hex_chars[bytes[i] >> 4];
        hex[i * 2 + 1] = hex_chars[bytes[i] & 0x0F];
    }
    hex[len * 2] = '\0';
}

// 将十六进制字符串转换为字节数组
int crypto_hex_to_bytes(const char *hex, uint8_t *bytes, size_t *len) {
    if (!hex || !bytes || !len) return CRYPTO_ERROR_INVALID_PARAM;
    
    size_t hex_len = strlen(hex);
    if (hex_len % 2 != 0) return CRYPTO_ERROR_INVALID_PARAM;
    
    *len = hex_len / 2;
    
    for (size_t i = 0; i < *len; i++) {
        char high = hex[i * 2];
        char low = hex[i * 2 + 1];
        
        uint8_t high_val, low_val;
        
        if (high >= '0' && high <= '9') high_val = high - '0';
        else if (high >= 'a' && high <= 'f') high_val = high - 'a' + 10;
        else if (high >= 'A' && high <= 'F') high_val = high - 'A' + 10;
        else return CRYPTO_ERROR_INVALID_PARAM;
        
        if (low >= '0' && low <= '9') low_val = low - '0';
        else if (low >= 'a' && low <= 'f') low_val = low - 'a' + 10;
        else if (low >= 'A' && low <= 'F') low_val = low - 'A' + 10;
        else return CRYPTO_ERROR_INVALID_PARAM;
        
        bytes[i] = (high_val << 4) | low_val;
    }
    
    return CRYPTO_SUCCESS;
} 